#include <iostream>
#include <fstream>
#include <ctime>
#include <cstdlib>
#include <cmath>
#include "anyoption.h"
#include "sbml/SBMLTypes.h"
#include "../libBioModel/Network.h"
#include "../libBioModel/Stoichiometry.h"
#include "../libBioModel/State.h"
#include "../libBioModel/InitialConds.h"

using namespace std;


/*
 * An implementation of Gillespie's Stochastic Simulation Algorith to solve 
 * the Chemical Master Equation based on Monte Carlo simulation.
 *
 * This program implements the direct method of SSA.
 *
 * Sep 22, 2009
 * Youfang Cao
 *
 */

int main(int argc, char* argv[])
{
	/* Variable definitions */
	char*		sbmlfile = 0;        	// SBML filename.
	char*		initfile = 0;        	// Inital condition filename.
	unsigned int	n_trajectories = 1;		// The total number of simulations to fire.
	unsigned int	n_timesteps = 100000;    		// The number of time steps.
	double		totaltime = 1000;            	// Total time for simulation.
	unsigned int	bpoint = 15000; 	// Starting point of numerarical integral.
	double		tol = 1E-10;       	// Tolerace for iteration stop.
	unsigned int	scalex = 101; 		// X scale of the landscape.
	unsigned int	scaley = 101; 		// Y scale of the landscape.
	unsigned int	ConvTimes = 5;    	// Times for error be less than tol, for which we can think the iteration for the landscape converged.
	int		verbose = 0;
	int		flag_landscape = 1;

	/* Variables for landscape */
	string		xid;
	string		yid;
	int		xindex = 0;
	int		yindex = 0;
	int		LandDim = 2;
	unsigned int*	Grid = 0;
	unsigned int*	dGrid = 0;
	double*		ProbLand = 0;
	unsigned long	ncounter = 0;
	unsigned long	dncounter = 0;
	char*		landfile = 0;
	unsigned int	xcn = 0;
	unsigned int	ycn = 0;
	double		err = 0;
	int		cconvtimes = 0;

	// Iterative variables
	unsigned int i=0, j=0, k=0; // some index variables.
	int f = 0;
	unsigned int r = 0;
	double t = 0;
	double dt = 0;

	// Variable definitions for Gillespie Algorithm
	InitialConds X0; // The initial state.
	State* X = 0; // The state.
	Network nw;
	unsigned int n_reactions = 0;
	double tau1 = 0;
	double tau2 = 0;
	double* rates = 0;
	Reaction* re;
	KineticLaw* klaw;
	double a0 = 0;
	double cum = 0;
	double cri = 0;

        /* CREATE AN OBJECT TO HANDLE COMMAND LINE */
        AnyOption *opt = new AnyOption();

        /* SET THE USAGE/HELP   */
        opt->addUsage( "" );
        opt->addUsage( "Usage: " );
        opt->addUsage( "" );
        opt->addUsage( " -h  --help  		Prints this help " );
        opt->addUsage( " -l  --land  		Landscapes or single simulations " );
        opt->addUsage( " -m  --sbml  		Specify the SBML file" );
        opt->addUsage( " -i  --init  		Specify the initial condition" );
        opt->addUsage( " -T  --ntrajectories 10	Number of simulations for a landscape " );
        opt->addUsage( " -S  --nsteps 20000     Number of time steps for each simulation " );
        opt->addUsage( " -t  --ttime 1000  	Total time for simulation " );
        opt->addUsage( " -b  --begin 15000  	The begining step for numerical integrals " );
        opt->addUsage( " -e  --tol 1E-8  	Error tolerace of landscape difference for iteration stop " );
        opt->addUsage( " -c  --conv 5  		Min consecutive number of times of err less than tol " );
        opt->addUsage( " -v  --verbose 0  	Switch for more output " );
        opt->addUsage( "" );
        opt->addUsage( "  Landscape parameters:" );
        opt->addUsage( " -D  --dim 2  		Dimension of landscape, 1 or 2(default) " );
        opt->addUsage( " -x  --scalex 101  	X scale of the 2D landscape " );
        opt->addUsage( " -y  --scaley 101  	Y scale of the 2D landscape " );
        opt->addUsage( " -X  --xid X 		X ID in SBML model " );
        opt->addUsage( " -Y  --yid Y 		Y ID in SBML model " );
        opt->addUsage( " -o  --ofile 		Landscape output file " );
        opt->addUsage( "" );

        /* SET THE OPTION STRINGS/CHARACTERS */
        opt->setFlag(  "help", 'h' );
        opt->setFlag(  "land", 'l' );
        opt->setOption(  "sbml", 'm' );
        opt->setOption(  "init", 'i' );
        opt->setOption(  "ntrajectories", 'T' );
        opt->setOption(  "nsteps", 'S' );
        opt->setOption(  "ttime", 't' );
        opt->setOption(  "begin", 'b' );
        opt->setOption(  "dim", 'D' );
        opt->setOption(  "scalex", 'x' );
        opt->setOption(  "scaley", 'y' );
        opt->setOption(  "xid", 'X' );
        opt->setOption(  "yid", 'Y' );
        opt->setOption(  "tol", 'e' );
        opt->setOption(  "conv", 'c' );
        opt->setOption(  "verbose", 'v' );
        opt->setOption(  "ofile", 'o' );

	/* go through the command line and get the options  */
        opt->processCommandArgs( argc, argv );

        /* GET THE VALUES */
        if( opt->getFlag( "help" ) || opt->getFlag( 'h' ) ) 
	{
                opt->printUsage();
		delete opt;
		return 1;
	}
        if( opt->getFlag( "land" ) || opt->getFlag( 'l' ) ) 
	{
		flag_landscape = 1;
	}
	else
	{
		flag_landscape = 0;
	}
	if( opt->getValue( 'm' ) != NULL  || opt->getValue( "sbml" ) != NULL  )
		sbmlfile = opt->getValue( 'm' );
	else
	{
		printf ("Please specify the SBML filename.\n\n");
		opt->printUsage();
		return 1;
	}
	if( opt->getValue( 'i' ) != NULL  || opt->getValue( "init" ) != NULL  )
		initfile = opt->getValue( 'i' );
	else
	{
		printf ("Please specify the initial condition.\n\n");
		opt->printUsage();
		return 1;
	}
	if( opt->getValue( 'T' ) != NULL  || opt->getValue( "ntrajectories" ) != NULL  )
		n_trajectories = atoi(opt->getValue( 'T' ));
	if( opt->getValue( 'S' ) != NULL  || opt->getValue( "nsteps" ) != NULL  )
		n_timesteps = atoi(opt->getValue( 'S' ));
	if( opt->getValue( 't' ) != NULL  || opt->getValue( "ttime" ) != NULL  )
		totaltime = atof(opt->getValue( 't' ));
	if( opt->getValue( 'b' ) != NULL  || opt->getValue( "begin" ) != NULL  )
		bpoint = atoi(opt->getValue( 'b' ));
	if( opt->getValue( 'D' ) != NULL  || opt->getValue( "dim" ) != NULL  )
		LandDim = atoi(opt->getValue( 'D' ));
	if( opt->getValue( 'x' ) != NULL  || opt->getValue( "scalex" ) != NULL  )
		scalex = atoi(opt->getValue( 'x' ));
	if( opt->getValue( 'y' ) != NULL  || opt->getValue( "scaley" ) != NULL  )
		scaley = atoi(opt->getValue( 'y' ));
	if( opt->getValue( 'X' ) != NULL  || opt->getValue( "xid" ) != NULL  )
		xid = opt->getValue( 'X' );
	if( opt->getValue( 'Y' ) != NULL  || opt->getValue( "yid" ) != NULL  )
		yid = opt->getValue( 'Y' );
	if( opt->getValue( 'e' ) != NULL  || opt->getValue( "tol" ) != NULL  )
		tol = atof(opt->getValue( 'e' ));
	if( opt->getValue( 'c' ) != NULL  || opt->getValue( "conv" ) != NULL  )
		ConvTimes = atoi(opt->getValue( 'c' ));
	if( opt->getValue( 'v' ) != NULL  || opt->getValue( "verbose" ) != NULL  )
		verbose = atoi(opt->getValue( 'v' ));
	if( opt->getValue( 'o' ) != NULL  || opt->getValue( "ofile" ) != NULL  )
		landfile = opt->getValue( 'o' );
	cout << endl;

	/* Initialize variables. */

	/* READ SBML MODEL */
	nw.d = readSBML(sbmlfile);
	nw.m = nw.d->getModel();

	//errors = nw.d->getNumErrors();
	
	/* read stoichiometic matrix */
	nw.stoichiometry = new Stoichiometry();
	f = nw.stoichiometry->genStoichiometricMatrix(nw.m);
	if (f == -1)
	{
		printf ("Exit caused by error.\n");
		return -1;
	}
	if (verbose > 0)
		nw.stoichiometry->printStoichiometricMatrix();
	
	//modelid = nw.m->getId();
	nw.n_compartments = nw.m->getNumCompartments();
	nw.n_species = nw.m->getNumSpecies();
	nw.n_rules = nw.m->getNumRules();
	n_reactions = nw.stoichiometry->getNumReactions();
	nw.n_reactions = nw.m->getNumReactions();

	/* read reaction IDs */
	nw.ReactionsID = new string[n_reactions];
	j = 0;
	for (i = 0; i < nw.n_reactions; i++)
	{
		nw.ReactionsID[j] = nw.m->getReaction(i)->getId();
		//cout << "heloo. " << n_reactions << " " << i << " " << j << " " << nw.ReactionsID[j] << endl;
		j++;
		if (nw.m->getReaction(i)->getReversible())
		{
			nw.ReactionsID[j] = nw.m->getReaction(i)->getId();
			j++;
		}
	}

	//cout << "heloo. " << endl;
	/* read compartment IDs */
	nw.CompartmentsID = new string[nw.n_compartments];
	for (i = 0; i < nw.n_compartments; i++)
	{
		nw.CompartmentsID[i] = nw.m->getCompartment(i)->getId();
	}

	/* read species IDs */
	nw.SpeciesID = new string[nw.n_species];
	if (verbose > 0)
	{
		printf ("Species:\n");
		printf ("id\tinitAmt\tinitCon\tconst\tboundcond\n");
	}
	for(i = 0; i < nw.n_species; i++)
	{
		/* Store id data in the array. */
		nw.SpeciesID[i] = nw.m->getSpecies(i)->getId();
	}
	
	/*
	 * Get rules ID list.
	 */
	if(verbose > 0)
	{
		printf ("Rules:\n");
	}
	nw.RulesID = new string[nw.n_rules];
	nw.n_assignmentRules = 0;
	for(i = 0; i < nw.n_rules; i++)
	{
		if (SBase_getTypeCode((SBase_t*)nw.m->getRule(i)) != SBML_ASSIGNMENT_RULE)
			continue;
			
		/* Store id data in the array. */
		nw.RulesID[i] = ((AssignmentRule*)nw.m->getRule(i))->getVariable();
		
		nw.n_assignmentRules ++;
	}

	/*
	 * Get parameters ID list.
	 */
	if(verbose > 0)
	{
		printf ("Global Parameters:\n");
	}
	nw.n_parameters = nw.m->getNumParameters();
	nw.ParametersID = new string[nw.n_parameters];
	for (i = 0; i < nw.n_parameters; i++)
	{
		nw.ParametersID[i] = nw.m->getParameter(i)->getId();
		if (verbose > 0)
		{
			printf ("%s\t%f\n", nw.ParametersID[i].data(), nw.m->getParameter(i)->getValue());
		}
	}

	/* Local parameter IDs */
	nw.LocalParamIDs = new string[20];

	/* Initialize rates */
	n_reactions = nw.stoichiometry->getNumReactions();
	rates = new double[n_reactions];
	for (i=0; i<n_reactions; i++)
	{
		rates[i] = 0;
	}

	/* READ INITIAL CONDITION */
	if (X0.readInitialConds(initfile) == -1)
	{
		printf ("ERROR:  Read initial condition file error.\n");
		return -1;
	}

	/* X and Y labels */
	xindex = nw.getSpeciesIndex(xid);
	yindex = nw.getSpeciesIndex(yid);
	//cout << xid << " " << xindex << " " << yid << " " << yindex << endl;
	if (xindex < 0)
	{
		cout << "Species X:" << xid << " does not exist." << endl;
		return 0;
	}
	if (LandDim == 2 && yindex < 0)
	{
		cout << "Species Y:" << yid << " does not exist." << endl;
		return 0;
	}

	/* Initialize Grid and ProbLand */
	if (flag_landscape == 1)
	{
		if (LandDim == 1)
		{
			Grid = new unsigned int[scalex];
			dGrid = new unsigned int[scalex];
			ProbLand = new double[scalex];
			for (i=0;i<scalex;i++)
			{
				Grid[i] = 0;
				dGrid[i] = 0;
				ProbLand[i] = 0.0;
			}
		}
		else
		{
			Grid = new unsigned int[scalex*scaley];
			dGrid = new unsigned int[scalex*scaley];
			ProbLand = new double[scalex*scaley];
			for (i=0;i<scalex;i++)
			{
				for (j=0;j<scaley;j++)
				{
					Grid[i*scaley+j] = 0;
					dGrid[i*scaley+j] = 0;
					ProbLand[i*scaley+j] = 0.0;
				}
			}
		}
	}

	/* START SSA ITERATIONS */
	srand(time(NULL));
	i = 0;
	X = new State();
	while (i < n_timesteps)
	{
		cout << "Run: " << i+1;
		X->CopyfromState(X0.getInitialState(0));
		if (verbose > 0) { X->PrintState(); }
		if (LandDim == 1)
		{
			for (k=0;k<scalex;k++)
			{
				dGrid[k] = 0;
			}
		}
		else
		{
			for (k=0;k<scalex;k++)
			{
				for (j=0;j<scaley;j++)
				{
					dGrid[k*scaley+j] = 0;
				}
			}
		}
		dncounter = 0;
		j = 0;
		t = 0;
		while (j < n_trajectories)// && t < totaltime)
		{
			/* Calculate transition rates for all reactions in current state X */
			a0 = 0;
			for (k=0; k<n_reactions; k++)
			{
				re = nw.m->getReaction(abs(nw.stoichiometry->getReactionAttribute(k)) - 1);
				klaw = re->getKineticLaw();
				rates[k] = nw.getReactionRate(klaw, k, X) * (double)(nw.stoichiometry->canReact(k, X));
				a0 = a0 + rates[k];
				//cout << rates[k] << " ";
			}
			//cout << endl;

			/* Find the tau1 and tau2 */
			tau1 = (double)rand()/(double)RAND_MAX;
			tau2 = (double)rand()/(double)RAND_MAX;

			/* Now, I got the delta t */
			dt = log(1/tau1)/a0;

			//cout << dt << " " << tau1 << " " << a0 << " " << tau2 << " " << a0*tau2 << endl;

			cum = 0;
			cri = tau2*a0;
			for (k=0; k<n_reactions; k++)
			{
				cum += rates[k];
				if (cum >= cri) { break; }
			}
			/* Now, I got which reaction to go */
			r = k;

			/* Update the system state X and time t */
			t = t + dt;
			nw.stoichiometry->updateStateWithReaction(r, X);

			if (verbose > 0) { X->PrintState(); }

			xcn = X->getStateValue(xindex);
			if (LandDim == 2) {
				ycn = X->getStateValue(yindex);
			}
			if (j > bpoint)
			{
				if (LandDim == 1)
				{
					if (xcn < scalex)
						Grid[xcn] += 1;
					dncounter ++;
				}
				else
				{
					if (xcn < scalex && ycn < scaley)
						Grid[xcn*scaley + ycn] += 1;
					dncounter ++;
				}
			}

			/* Now, Go to next iteration */
			j++;
		}
		/* Error check */
		err = 0;
		if (ncounter != 0)
		{
			for (k=0; k<scalex; k++)
			{
				if (LandDim == 1)
				{
					err += pow( ((double)dncounter * Grid[k] - (double)ncounter * dGrid[k]) / ((double)ncounter*((double)ncounter+(double)dncounter)), 2 );
	                                Grid[k] += dGrid[k];
				}
				else
				{
					for (j=0; j<scaley; j++)
					{
						err += pow( ((double)dncounter * Grid[k*scaley+j] - (double)ncounter * dGrid[k*scaley+j]) / ((double)ncounter*((double)ncounter+(double)dncounter)), 2 );
						Grid[k*scaley+j] += dGrid[k*scaley+j];
					}
				}
			}
			err = sqrt(err);
			ncounter += dncounter;
		}
		else
		{
			for (k=0; k<scalex; k++)
			{
				if (LandDim == 1)
				{
					Grid[k] += dGrid[k];
				}
				else
				{
					for (j=0; j<scaley; j++)
					{
						Grid[k*scaley+j] += dGrid[k*scaley+j];
					}
				}
			}
			ncounter += dncounter;
			err = 10;
		}

		cout << ", err=" << scientific << err << endl;
		//printf ("nStep=%d, err=%E\n", i, err);

		if (err < tol)
		{
			cconvtimes ++;
		}
		else
		{
			cconvtimes = 0;
		}


		/* Go ahead to next run */
		i ++;
	}

	if (flag_landscape == 1)
	{
		ofstream fland;
		fland.open(landfile);

		if (LandDim == 1)
		{
			fland << xid << " Prob" << endl; /* How clean is this? Why do I still need fprintf stuff? */
		}
		else
		{
			fland << xid << " " << yid << " Prob" << endl; /* How clean is this? Why do I still need fprintf stuff? */
		}

		for (i=0;i<scalex;i++)
		{
			if (LandDim == 1)
			{
				ProbLand[i] = (double)(Grid[i])/(double)ncounter;
				fland << i << " " << scientific << ProbLand[i] << endl;
			}
			else
			{
				for (j=0;j<scaley;j++)
				{
					ProbLand[i*scaley+j] = (double)(Grid[i*scaley+j])/(double)ncounter;
					fland << i << " " << j << " " << scientific << ProbLand[i*scaley+j] << endl;
				}
			}
		}

		fland.close();
	}

}


/*
 */




